import pygame
import numpy as np
import math
import random
from typing import List, Tuple, Dict, Any


class ArmEnv:
    # 常量定义
    SCREEN_WIDTH = 640
    SCREEN_HEIGHT = 640
    WORLD_WIDTH = 5.0  # 米
    WORLD_HEIGHT = 5.0  # 米
    SCALE = SCREEN_WIDTH / WORLD_WIDTH  # 像素/米

    # 颜色定义
    WHITE = (255, 255, 255)
    BLACK = (0, 0, 0)
    GRAY = (128, 128, 128)
    RED = (255, 0, 0)
    GREEN = (0, 255, 0)
    BLUE = (0, 120, 255)
    YELLOW = (255, 255, 0)
    PURPLE = (180, 0, 255)
    PINK = (255, 100, 180)
    ORANGE = (255, 150, 0)

    # 机械臂参数
    LINK_LENGTHS = [2.0, 1.6, 0.4]  # 三节长度

    LINK_ANGLE_ACC_LIMIT = 0.05

    def __init__(self):
        # 初始化pygame
        pygame.init()
        self.screen = pygame.display.set_mode((self.SCREEN_WIDTH, self.SCREEN_HEIGHT))
        pygame.display.set_caption("2D Robotic Arm Environment")
        self.clock = pygame.time.Clock()

        # 物理参数
        self.gravity = 9.8  # m/s² (y轴方向)
        self.dt = 0.1  # 时间步长

        # 机械臂状态
        self.base_pos = np.array(
            [self.WORLD_WIDTH / 2, self.WORLD_HEIGHT / 2]
        )  # 基座位置
        self.joint_angles = [0.0, 0.0, 0.0]  # 三个关节角度 (弧度)
        self.gripper_state = False  # 夹爪状态 (False=松开, True=抓取)

        # 目标关节角度（用于插值）
        self.target_joint_angles = self.joint_angles.copy()

        # 环境物体
        self.tables = []
        self.boxes = []

        # 重置环境
        self.reset()

    def reset(self):
        """重置环境到初始状态"""
        # 重置机械臂
        self.joint_angles = [0.0, 0.0, 0.0]
        self.target_joint_angles = self.joint_angles.copy()
        self.gripper_state = False

        # 清除现有物体
        self.tables = []
        self.boxes = []

        # 创建桌子 (3-5张)
        num_tables = random.randint(3, 5)
        table_colors = [self.GREEN, self.BLUE, self.YELLOW]

        for _ in range(num_tables):
            # 确保桌子不会重叠
            valid_position = False
            while not valid_position:
                x = random.uniform(0.5, self.WORLD_WIDTH - 0.5)
                y = random.uniform(0.5, self.WORLD_HEIGHT - 0.5)
                new_table = {
                    "pos": np.array([x, y]),
                    "size": np.array([0.5, 0.5]),
                    "color": random.choice(table_colors),
                }

                # 检查是否与现有桌子重叠
                overlap = False
                for table in self.tables:
                    if self._check_rect_overlap(
                        new_table["pos"], new_table["size"], table["pos"], table["size"]
                    ):
                        overlap = True
                        break

                if not overlap:
                    self.tables.append(new_table)
                    valid_position = True

        # 创建箱子 (4-6个)
        num_boxes = random.randint(4, 6)
        box_colors = [self.RED, self.PURPLE, self.PINK, self.ORANGE]

        for _ in range(num_boxes):
            # 箱子可以在地面或桌子上
            if self.tables and random.random() > 0.3:  # 70%几率在桌子上
                table = random.choice(self.tables)
                x = table["pos"][0] + random.uniform(-0.15, 0.15)
                y = table["pos"][1] - table["size"][1] / 2 - 0.1
            else:  # 在地面上
                x = random.uniform(0.2, self.WORLD_WIDTH - 0.2)
                y = self.WORLD_HEIGHT - 0.1  # 地面高度

            self.boxes.append(
                {
                    "pos": np.array([x, y]),
                    "size": np.array([0.2, 0.2]),
                    "color": random.choice(box_colors),
                    "velocity": np.array([0.0, 0.0]),
                    "on_table": False,
                    "grabbed": False,
                }
            )

    def step(self, action: np.ndarray):
        """
        执行动作并更新环境
        action: 三个关节的角度变化量 [Δθ0, Δθ1, Δθ2, gripper_action]
                其中 gripper_action: 0=无动作, 1=抓取, 2=松开
        """
        # 更新目标关节角度
        for i in range(3):
            acc = np.clip(
                action[i], -self.LINK_ANGLE_ACC_LIMIT, self.LINK_ANGLE_ACC_LIMIT
            )
            self.joint_angles[i] += acc

        # 处理夹爪动作
        if action[3] == 1:  # 抓取
            self.gripper_state = True
        elif action[3] == 2:  # 松开
            self.gripper_state = False

        # 更新被抓取的箱子
        self._update_grabbed_box()

        # 更新物理模拟
        self._physics_update()

        # 获取观测值
        observation = self._get_observation()

        # 渲染环境
        self.render()

        return observation

    def _update_grabbed_box(self):
        """更新被抓取的箱子状态"""
        # 获取末端执行器位置
        end_effector_pos = self._forward_kinematics()

        # 如果夹爪是闭合的，检查是否有箱子被抓取
        if self.gripper_state:
            for box in self.boxes:
                if not box["grabbed"]:
                    # 计算末端执行器与箱子的距离
                    distance = np.linalg.norm(end_effector_pos - box["pos"])

                    # 如果距离足够近且没有其他箱子被抓取
                    if distance < 0.15 and not any(b["grabbed"] for b in self.boxes):
                        box["grabbed"] = True
                        box["velocity"] = np.array([0.0, 0.0])  # 重置速度
        else:
            # 松开所有被抓取的箱子
            for box in self.boxes:
                if box["grabbed"]:
                    box["grabbed"] = False

    def _physics_update(self):
        """更新物理模拟（重力和碰撞）"""
        for box in self.boxes:
            if box["grabbed"]:
                # 被抓取的箱子跟随末端执行器 - 箱子中心与末端重合
                end_effector_pos = self._forward_kinematics()
                box["pos"] = end_effector_pos  # 修改这里：箱子中心与末端执行器位置重合
                box["velocity"] = np.array([0.0, 0.0])
                box["on_table"] = False
            else:
                # 应用重力
                box["velocity"][1] += self.gravity * self.dt

                # 更新位置
                new_pos = box["pos"] + box["velocity"] * self.dt

                # 检查与地面的碰撞
                ground_level = self.WORLD_HEIGHT - box["size"][1] / 2
                if new_pos[1] > ground_level:
                    new_pos[1] = ground_level
                    box["velocity"][1] = 0.0
                    box["on_table"] = False

                # 检查与桌子的碰撞
                box["on_table"] = False
                for table in self.tables:
                    table_top = table["pos"][1] - table["size"][1] / 2
                    table_bottom = table["pos"][1] + table["size"][1] / 2

                    # 检查箱子是否在桌面上方
                    if (
                        new_pos[0] > table["pos"][0] - table["size"][0] / 2
                        and new_pos[0] < table["pos"][0] + table["size"][0] / 2
                        and new_pos[1] > table_top - box["size"][1] / 2
                        and new_pos[1] < table_bottom + box["size"][1] / 2
                    ):

                        # 碰撞响应
                        new_pos[1] = table_top - box["size"][1] / 2
                        box["velocity"][1] = 0.0
                        box["on_table"] = True

                # 更新位置
                box["pos"] = new_pos

    def _forward_kinematics(self) -> np.ndarray:
        """计算机械臂末端执行器位置"""
        x = self.base_pos[0]
        y = self.base_pos[1]

        # 计算每个关节的位置
        angle = self.joint_angles[0]
        x += self.LINK_LENGTHS[0] * np.cos(angle)
        y += self.LINK_LENGTHS[0] * np.sin(angle)

        angle += self.joint_angles[1]
        x += self.LINK_LENGTHS[1] * np.cos(angle)
        y += self.LINK_LENGTHS[1] * np.sin(angle)

        angle += self.joint_angles[2]
        x += self.LINK_LENGTHS[2] * np.cos(angle)
        y += self.LINK_LENGTHS[2] * np.sin(angle)

        return np.array([x, y])

    def _get_observation(self) -> Dict[str, Any]:
        """获取环境观测值"""
        # 1. 获取RGB图像
        self.render()
        rgb_array = pygame.surfarray.array3d(self.screen)
        # 转换维度顺序为 (C, H, W)
        rgb_array = np.transpose(rgb_array, (2, 0, 1))

        # 2. 获取物体状态信息
        objects = []

        # 添加桌子信息
        for table in self.tables:
            objects.append(
                [
                    "table",
                    table["pos"][0],
                    table["pos"][1],
                    0.0,  # sin(rotation)
                    1.0,  # cos(rotation)
                    table["color"],
                ]
            )

        # 添加箱子信息
        for box in self.boxes:
            objects.append(
                [
                    "box",
                    box["pos"][0],
                    box["pos"][1],
                    0.0,  # sin(rotation)
                    1.0,  # cos(rotation)
                    box["color"],
                ]
            )

        # 3. 获取机械臂状态
        arm_state = {
            "base_pos": self.base_pos.tolist(),
            "joint_angles": self.joint_angles.copy(),
            "end_effector_pos": self._forward_kinematics().tolist(),
            "gripper_state": self.gripper_state,
        }

        return {"rgb_image": rgb_array, "objects": objects, "arm_state": arm_state}

    def render(self):
        """渲染环境"""
        self.screen.fill(self.WHITE)

        # 绘制墙壁
        pygame.draw.rect(self.screen, self.BLACK, (0, 0, self.SCREEN_WIDTH, 5))
        pygame.draw.rect(self.screen, self.BLACK, (0, 0, 5, self.SCREEN_HEIGHT))
        pygame.draw.rect(
            self.screen, self.BLACK, (self.SCREEN_WIDTH - 5, 0, 5, self.SCREEN_HEIGHT)
        )
        pygame.draw.rect(
            self.screen, self.BLACK, (0, self.SCREEN_HEIGHT - 5, self.SCREEN_WIDTH, 5)
        )

        # 绘制桌子
        for table in self.tables:
            x, y = self._world_to_screen(table["pos"])
            w, h = table["size"] * self.SCALE
            pygame.draw.rect(self.screen, table["color"], (x - w / 2, y - h / 2, w, h))

        # 绘制箱子
        for box in self.boxes:
            x, y = self._world_to_screen(box["pos"])
            w, h = box["size"] * self.SCALE
            pygame.draw.rect(self.screen, box["color"], (x - w / 2, y - h / 2, w, h))

        # 绘制机械臂
        self._draw_arm()

        pygame.display.flip()
        self.clock.tick(30)  # 30 FPS

    def _draw_arm(self):
        """绘制机械臂"""
        # 基座位置
        base_screen = self._world_to_screen(self.base_pos)

        # 计算各关节位置
        positions = [self.base_pos]
        angle = self.joint_angles[0]
        pos1 = self.base_pos + np.array(
            [self.LINK_LENGTHS[0] * np.cos(angle), self.LINK_LENGTHS[0] * np.sin(angle)]
        )
        positions.append(pos1)

        angle += self.joint_angles[1]
        pos2 = pos1 + np.array(
            [self.LINK_LENGTHS[1] * np.cos(angle), self.LINK_LENGTHS[1] * np.sin(angle)]
        )
        positions.append(pos2)

        angle += self.joint_angles[2]
        end_effector = pos2 + np.array(
            [self.LINK_LENGTHS[2] * np.cos(angle), self.LINK_LENGTHS[2] * np.sin(angle)]
        )
        positions.append(end_effector)

        # 转换为屏幕坐标
        screen_positions = [self._world_to_screen(p) for p in positions]

        # 绘制连杆
        pygame.draw.line(
            self.screen, self.BLACK, screen_positions[0], screen_positions[1], 8
        )
        pygame.draw.line(
            self.screen, self.GRAY, screen_positions[1], screen_positions[2], 6
        )
        pygame.draw.line(
            self.screen, self.GRAY, screen_positions[2], screen_positions[3], 4
        )

        # 绘制关节
        for i, pos in enumerate(screen_positions[:3]):
            pygame.draw.circle(self.screen, self.BLACK, pos, 10)

        # 绘制末端执行器
        if self.gripper_state:  # 抓取状态
            pygame.draw.circle(self.screen, self.BLACK, screen_positions[3], 8)
        else:  # 松开状态
            pygame.draw.circle(self.screen, self.BLACK, screen_positions[3], 8, 2)

    def _world_to_screen(self, world_pos: np.ndarray) -> Tuple[int, int]:
        """世界坐标转换为屏幕坐标"""
        x = int(world_pos[0] * self.SCALE)
        y = int(world_pos[1] * self.SCALE)
        return (x, y)

    def _wrap_angle(self, angle: float) -> float:
        """将角度限制在[-π, π]范围内"""
        while angle > np.pi:
            angle -= 2 * np.pi
        while angle < -np.pi:
            angle += 2 * np.pi
        return angle

    def _check_rect_overlap(self, pos1, size1, pos2, size2):
        """检查两个矩形是否重叠"""
        return (abs(pos1[0] - pos2[0]) < (size1[0] / 2 + size2[0] / 2)) and (
            abs(pos1[1] - pos2[1]) < (size1[1] / 2 + size2[1] / 2)
        )

    def close(self):
        """关闭环境"""
        pygame.quit()


# 测试环境
if __name__ == "__main__":
    env = ArmEnv()
    running = True

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        # 随机动作测试
        action = np.array(
            [
                random.uniform(-0.1, 0.1),  # Δθ0
                random.uniform(-0.1, 0.1),  # Δθ1
                random.uniform(-0.1, 0.1),  # Δθ2
                random.choice([0, 1, 2]),  # 夹爪动作
            ]
        )

        env.step(action)

    env.close()